# json_to_docx_fixed_encoding.py
import json
import argparse
import sys
from pathlib import Path
from docx import Document
from docx.shared import Pt, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH
import traceback

class JsonToDocxConverter:
    def __init__(self):
        self.doc = None

    def convert_json_to_docx(self, json_data, docx_path):
        """将 JSON 数据转换为 DOCX 文档 - 修复编码版本"""
        try:
            print("正在创建 DOCX 文档...")
            self.doc = Document()

            # 应用元数据
            self._apply_metadata(json_data.get('metadata', {}))

            # 应用内容
            self._apply_content(json_data.get('content', []))
            self._apply_paragraphs(json_data.get('paragraphs', []))
            self._apply_tables(json_data.get('tables', []))

            # 保存文档
            print(f"正在保存文档: {docx_path}")
            self.doc.save(str(docx_path))
            return {
                "success": True,
                "output_path": str(docx_path),
                "error": None
            }

        except Exception as e:
            error_msg = f"转换失败: {str(e)}"
            print(f"错误详情: {error_msg}")
            traceback.print_exc()
            return {
                "success": False,
                "output_path": None,
                "error": error_msg
            }

    def _safe_decode(self, text):
        """安全解码文本，处理中文"""
        if text is None:
            return ""
        try:
            # 如果是字符串，直接返回
            if isinstance(text, str):
                return text
            # 如果是其他类型，转换为字符串
            else:
                return str(text)
        except:
            return ""

    def _apply_metadata(self, metadata):
        """应用文档元数据 - 修复编码"""
        try:
            if metadata:
                title = self._safe_decode(metadata.get('title'))
                author = self._safe_decode(metadata.get('author'))

                if title:
                    self.doc.core_properties.title = title
                if author:
                    self.doc.core_properties.author = author

                print("元数据应用完成")
        except Exception as e:
            print(f"应用元数据时出错: {e}")

    def _apply_content(self, content):
        """应用详细内容 - 修复编码"""
        if not content:
            return

        print(f"正在处理 {len(content)} 个内容项...")
        for i, item in enumerate(content):
            try:
                if item.get('type') == 'paragraph':
                    self._add_formatted_paragraph(item)
                elif item.get('text'):
                    paragraph = self.doc.add_paragraph()
                    text = self._safe_decode(item['text'])
                    paragraph.add_run(text)
            except Exception as e:
                print(f"处理内容项 {i} 时出错: {e}")
                if item.get('text'):
                    paragraph = self.doc.add_paragraph()
                    text = self._safe_decode(item['text'])
                    paragraph.add_run(text)

    def _apply_paragraphs(self, paragraphs):
        """应用简化版本的段落 - 修复编码"""
        if not paragraphs:
            return

        print(f"正在处理 {len(paragraphs)} 个简化段落...")
        for i, para_data in enumerate(paragraphs):
            try:
                paragraph = self.doc.add_paragraph()
                text = self._safe_decode(para_data.get('text', ''))
                paragraph.add_run(text)
            except Exception as e:
                print(f"处理简化段落 {i} 时出错: {e}")

    def _add_formatted_paragraph(self, para_data):
        """添加带格式的段落 - 修复编码"""
        paragraph = self.doc.add_paragraph()

        # 设置段落样式
        if para_data.get('style'):
            try:
                style_name = self._safe_decode(para_data['style'])
                paragraph.style = style_name
            except Exception as e:
                print(f"设置样式时出错: {e}")

        # 设置对齐方式
        alignment = para_data.get('alignment', 'left')
        alignment_map = {
            'left': WD_ALIGN_PARAGRAPH.LEFT,
            'center': WD_ALIGN_PARAGRAPH.CENTER,
            'right': WD_ALIGN_PARAGRAPH.RIGHT,
            'justify': WD_ALIGN_PARAGRAPH.JUSTIFY
        }
        paragraph.alignment = alignment_map.get(alignment, WD_ALIGN_PARAGRAPH.LEFT)

        # 添加运行
        if para_data.get('runs'):
            for run_data in para_data['runs']:
                self._add_formatted_run(paragraph, run_data)
        elif para_data.get('text'):
            text = self._safe_decode(para_data['text'])
            paragraph.add_run(text)

    def _add_formatted_run(self, paragraph, run_data):
        """添加带格式的文本运行 - 修复编码"""
        run_text = self._safe_decode(run_data.get('text', ''))
        if not run_text:
            return

        run = paragraph.add_run(run_text)

        # 应用文本格式
        if run_data.get('bold'):
            run.bold = bool(run_data['bold'])
        if run_data.get('italic'):
            run.italic = bool(run_data['italic'])
        if run_data.get('underline'):
            run.underline = bool(run_data['underline'])

    def _apply_tables(self, tables):
        """应用表格 - 修复编码"""
        if not tables:
            return

        print(f"正在处理 {len(tables)} 个表格...")
        for table_idx, table_data in enumerate(tables):
            try:
                rows_data = table_data.get('rows', [])
                if not rows_data:
                    continue

                # 计算最大列数
                max_cols = 0
                for row in rows_data:
                    if isinstance(row, list):
                        max_cols = max(max_cols, len(row))
                    else:
                        cells = [cell for cell in row if isinstance(cell, dict)]
                        max_cols = max(max_cols, len(cells))

                if max_cols == 0:
                    continue

                # 创建表格
                num_rows = len(rows_data)
                table = self.doc.add_table(rows=num_rows, cols=max_cols)

                # 填充表格内容
                for row_idx, row_data in enumerate(rows_data):
                    if isinstance(row_data, list):
                        # 简单列表格式
                        for col_idx, cell_text in enumerate(row_data):
                            if row_idx < num_rows and col_idx < max_cols:
                                text = self._safe_decode(cell_text)
                                table.cell(row_idx, col_idx).text = text
                    else:
                        # 单元格对象格式
                        for cell_data in row_data:
                            if isinstance(cell_data, dict):
                                row_idx = cell_data.get('row', row_idx)
                                col_idx = cell_data.get('column', cell_data.get('col', 0))
                                cell_text = self._safe_decode(cell_data.get('text', ''))
                                if row_idx < num_rows and col_idx < max_cols:
                                    table.cell(row_idx, col_idx).text = cell_text

                print(f"表格 {table_idx} 创建完成")

            except Exception as e:
                print(f"创建表格 {table_idx} 时出错: {e}")

def load_json_input(input_source):
    """加载 JSON 数据 - 修复编码"""
    try:
        if isinstance(input_source, str):
            # 尝试作为文件路径处理
            file_path = Path(input_source)
            if file_path.exists() and file_path.suffix.lower() in ['.json', '.txt']:
                print(f"从文件加载 JSON: {input_source}")
                # 确保使用 UTF-8 编码读取
                with open(file_path, 'r', encoding='utf-8') as f:
                    return json.load(f)

            # 尝试作为 JSON 字符串解析
            try:
                print("尝试解析 JSON 字符串...")
                return json.loads(input_source)
            except json.JSONDecodeError:
                raise ValueError("输入不是有效的 JSON 文件路径或 JSON 字符串")

        elif isinstance(input_source, dict):
            return input_source

        else:
            raise ValueError("不支持的输入类型")

    except Exception as e:
        raise ValueError(f"加载 JSON 数据失败: {str(e)}")

def main():
    # 设置控制台编码为 UTF-8
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

    parser = argparse.ArgumentParser(description='JSON 转 DOCX 转换工具 - 修复编码版本')

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--file', '-f', help='输入的 JSON 文件路径')
    group.add_argument('--json', '-j', help='直接的 JSON 数据字符串')

    parser.add_argument('--output', '-o', help='输出的 DOCX 文件路径')

    args = parser.parse_args()

    converter = JsonToDocxConverter()

    try:
        # 确定输入源
        if args.file:
            input_source = args.file
        else:
            input_source = args.json

        # 加载 JSON 数据
        print("正在加载 JSON 数据...")
        json_data = load_json_input(input_source)
        print("JSON 数据加载成功")

        # 确定输出路径
        if args.output:
            output_path = Path(args.output)
        elif args.file:
            input_path = Path(args.file)
            output_path = input_path.with_suffix('.docx')
        else:
            output_path = Path("output.docx")

        print(f"输出文件: {output_path}")

        # 执行转换
        result = converter.convert_json_to_docx(json_data, output_path)

        # 输出结果
        print(json.dumps(result, ensure_ascii=False))

        if result['success']:
            print("✅ 转换成功！")
            sys.exit(0)
        else:
            print("❌ 转换失败")
            sys.exit(1)

    except Exception as e:
        error_result = {
            "success": False,
            "output_path": None,
            "error": f"程序执行错误: {str(e)}"
        }
        print(json.dumps(error_result, ensure_ascii=False))
        sys.exit(1)

if __name__ == "__main__":
    main()